/*
 * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlin.generators.tree.printer

import org.jetbrains.kotlin.generators.tree.*
import org.jetbrains.kotlin.generators.tree.imports.ImportCollecting
import org.jetbrains.kotlin.generators.tree.imports.ImportCollector
import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil
import org.jetbrains.kotlin.utils.IndentingPrinter
import org.jetbrains.kotlin.utils.SmartPrinter
import java.io.File

private val COPYRIGHT by lazy { File("license/COPYRIGHT_HEADER.txt").readText() }

class GeneratedFile(val file: File, val newText: String)

private fun getPathForFile(generationPath: File, packageName: String, typeName: String): File {
    val dir = generationPath.resolve(packageName.replace(".", "/"))
    return File(dir, "$typeName.kt")
}

private data class PrinterAndImportCollector(
    val printer: SmartPrinter,
    val importCollector: ImportCollecting,
) : IndentingPrinter by printer,
    ImportCollecting by importCollector,
    ImportCollectingPrinter

fun ImportCollectingPrinter.withNewPrinter(printer: SmartPrinter, body: ImportCollectingPrinter.() -> Unit) {
    PrinterAndImportCollector(printer, this).apply(body)
}

fun <GeneratedType, TypePrinter> printGeneratedTypesIntoSingleFile(
    generatedTypes: List<GeneratedType>,
    generationPath: File,
    treeGeneratorReadMe: String,
    packageName: String,
    fileNameWithoutExtension: String,
    fileSuppressions: List<String> = emptyList(),
    makeTypePrinter: (ImportCollectingPrinter) -> TypePrinter,
    printType: TypePrinter.(GeneratedType) -> Unit,
): GeneratedFile {
    val stringBuilder = StringBuilder()
    val file = getPathForFile(generationPath, packageName, fileNameWithoutExtension)
    val importCollector = ImportCollector(packageName)
    val printer = PrinterAndImportCollector(SmartPrinter(stringBuilder), importCollector)

    val typePrinter = makeTypePrinter(printer)
    var isFirst = true
    for (generatedType in generatedTypes) {
        if (isFirst) {
            isFirst = false
        } else {
            printer.println()
        }
        typePrinter.printType(generatedType)
    }

    return GeneratedFile(
        file,
        buildString {
            appendLine(COPYRIGHT)
            appendLine()
            append(GeneratorsFileUtil.GENERATED_MESSAGE_PREFIX)
            append(treeGeneratorReadMe)
            appendLine(".")
            appendLine(GeneratorsFileUtil.GENERATED_MESSAGE_SUFFIX)
            appendLine()
            if (fileSuppressions.isNotEmpty()) {
                fileSuppressions.joinTo(this, prefix = "@file:Suppress(", postfix = ")\n\n") { "\"$it\"" }
            }
            appendLine("package $packageName")
            appendLine()
            if (importCollector.printAllImports(this)) {
                appendLine()
            }
            append(stringBuilder)
        }
    )
}

fun printGeneratedType(
    generationPath: File,
    treeGeneratorReadMe: String,
    packageName: String,
    typeName: String,
    fileSuppressions: List<String> = emptyList(),
    body: ImportCollectingPrinter.() -> Unit,
): GeneratedFile =
    printGeneratedTypesIntoSingleFile(
        listOf(null),
        generationPath,
        treeGeneratorReadMe,
        packageName,
        typeName,
        fileSuppressions,
        makeTypePrinter = { it },
        printType = { body() }
    )

/**
 * Used to clean up files in [generationPath] that were not generated by calls to [Context.generateTree].
 *
 * @param generationPath Where to put the generated files.
 * @param treeGeneratorReadme A relative path to the README file of the tree generator. This path is mentioned in the auto-generation
 * warning in each generated file.
 */
class TreeGenerator(private val generationPath: File, val treeGeneratorReadme: String) {
    val generatedFiles = mutableListOf<GeneratedFile>()

    fun run(body: TreeGenerator.() -> Unit) {
        body(this)
        val previouslyGeneratedFiles = GeneratorsFileUtil.collectPreviouslyGeneratedFiles(generationPath)
        generatedFiles.forEach { GeneratorsFileUtil.writeFileIfContentChanged(it.file, it.newText, logNotChanged = false) }
        GeneratorsFileUtil.removeExtraFilesFromPreviousGeneration(previouslyGeneratedFiles, generatedFiles.map { it.file })
    }

    fun <Element> printElements(
        model: Model<Element>,
        createElementPrinter: (ImportCollectingPrinter) -> AbstractElementPrinter<Element, *>,
    ) where Element : AbstractElement<Element, *, *> {
        val elementsToPrint = model.elements.filter { it.doPrint }
        elementsToPrint.mapTo(generatedFiles) { element ->
            printGeneratedType(
                generationPath,
                treeGeneratorReadme,
                element.packageName,
                element.typeName,
            ) { createElementPrinter(this).printElement(element) }
        }
    }

    fun <Element, ElementField, Implementation> printElementImplementations(
        implementations: List<Implementation>,
        createImplementationPrinter: (ImportCollectingPrinter) -> AbstractImplementationPrinter<Implementation, Element, ElementField>,
    ) where Element : AbstractElement<Element, ElementField, Implementation>,
            ElementField : AbstractField<ElementField>,
            Implementation : AbstractImplementation<Implementation, Element, ElementField> {
        val implementationsToPrint = implementations.filter { it.doPrint }
        implementationsToPrint.mapTo(generatedFiles) { implementation ->
            printGeneratedType(
                generationPath,
                treeGeneratorReadme,
                implementation.packageName,
                implementation.typeName,
                fileSuppressions = listOf("DuplicatedCode"),
            ) { createImplementationPrinter(this).printImplementation(implementation) }
        }
    }

    fun <Element, ElementField> printElementBuilders(
        builders: List<Builder<ElementField, Element>>,
        createBuilderPrinter: ((ImportCollectingPrinter) -> AbstractBuilderPrinter<Element, *, ElementField>),
    ) where Element : AbstractElement<Element, ElementField, *>,
            ElementField : AbstractField<ElementField> {
        builders.mapTo(generatedFiles) { builder ->
            printGeneratedType(
                generationPath,
                treeGeneratorReadme,
                builder.packageName,
                builder.typeName,
                fileSuppressions = listOf("DuplicatedCode", "unused"),
            ) {
                createBuilderPrinter(this).printBuilder(builder)
            }
        }
    }

    fun <Element, ElementField> printVisitors(
        model: Model<Element>,
        createVisitorPrinters: List<Pair<ClassRef<*>, (ImportCollectingPrinter, ClassRef<*>) -> AbstractVisitorPrinter<Element, ElementField>>>,
    ) where Element : AbstractElement<Element, ElementField, *>,
            ElementField : AbstractField<ElementField> {
        createVisitorPrinters.mapTo(generatedFiles) { (visitorClass, createVisitorPrinter) ->
            printGeneratedType(generationPath, treeGeneratorReadme, visitorClass.packageName, visitorClass.simpleName) {
                createVisitorPrinter(this, visitorClass).printVisitor(model.elements)
            }
        }
    }
}
